import torchvision
import torch

# train_data = torchvision.datasets.ImageNet('./data_image_net', split='train', download=True,
#                                            transform=torchvision.transforms.ToTensor())
from torch.nn import Linear

vgg16_false = torchvision.models.vgg16(pretrained=False)
# vgg16_true = torchvision.models.vgg16(pretrained=True)
train_data = torchvision.datasets.CIFAR10('./dataset', train=True, transform=torchvision.transforms.ToTensor(), download=True)

# 在已训练好的模型上添加线性
vgg16_false.add_module('add_linear', Linear(1000, 10))
vgg16_false.classifier[6] = Linear(4096, 10)

print(vgg16_false)
